#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/kprobes.h>
#include <linux/sched.h>
#include <linux/delay.h>
#include <linux/types.h>
#include <linux/blk_types.h>
#include <linux/bio.h>
#include <linux/blkdev.h>
#include <linux/blk-mq.h>
#include <linux/wait.h>
#include <linux/time.h>

static DECLARE_WAIT_QUEUE_HEAD(my_wait_queue);

// 延时时间，单位:ms
static int sleep_time = 10;
module_param(sleep_time, int, 0);

// 故障磁盘名，比如:sda
static char error_disk[20];
module_param_string(error_disk, error_disk, 15, 0);

static int kprobe_io_pre_handler(struct kprobe *p, struct pt_regs *regs)
{
    struct task_struct *current_task = current;

    if (current_task->flags & PF_KTHREAD)
    {
        printk(KERN_INFO "a kernel thread enter kprobe, current process is %s - %d", current_task->comm, current_task->pid);
        return 0;
    }
    
    printk(KERN_INFO "a user thread enter kprobe, current process is %s - %d, sleep time is %d",
           current_task->comm, current_task->pid, sleep_time);
        
    if (!error_disk[0]) {
        printk(KERN_INFO "no input error disk");
        return -1;
    }
    printk(KERN_INFO "input error disk is %s", error_disk);

    // 添加白名单，在白名单中的进程不进行sleep
    char *whilt_process[] = {"vim", "mount", "python", "rmmod", "insmod", NULL};
    int i = 0;
    for (i = 0; whilt_process[i] != NULL; i++)
    {
        if (strcmp(current_task->comm, whilt_process[i]) == 0) 
        {
            printk(KERN_INFO "whilt process : %s, no sleep", current_task->comm);
            return 0;
        }
    }

    // 添加阻塞名单，只阻塞名单中的进程
    char *block_process[] = {"fio", "dd", NULL};
    for (i = 0; block_process[i] != NULL; i++)
    {
        if (strcmp(current_task->comm, block_process[i]) == 0) 
        {
            printk(KERN_INFO "block process : %s, start to sleep ... ...", current_task->comm);
            break;
        }
    }

    DECLARE_WAITQUEUE(wait, current_task);
    set_current_state(TASK_UNINTERRUPTIBLE);
    add_wait_queue(&my_wait_queue, &wait);

    unsigned long timeout_jiffies = msecs_to_jiffies(sleep_time);
    schedule_timeout_uninterruptible(timeout_jiffies);

    set_current_state(TASK_RUNNING);
    remove_wait_queue(&my_wait_queue, &wait);

    printk(KERN_INFO "block process : %s, sleep ok", current_task->comm);
    return 0;
}

void kprobe_io_post_handler(struct kprobe *p, struct pt_regs *regs, unsigned long flags)
{
    struct task_struct *current_task = current;
    printk(KERN_INFO "current process is %s - %d, exit function %s", current_task->comm, current_task->pid, p->symbol_name);
}

static struct kprobe kp = {
    .symbol_name = "blk_mq_get_tag",
    .pre_handler = kprobe_io_pre_handler,
    .post_handler = kprobe_io_post_handler,
};

static int __init kprobe_io_init(void)
{
    int ret = register_kprobe(&kp);
    if (ret < 0)
        printk(KERN_ERR "Failed to register kprobe");
    return 0;
}

static void __exit kprobe_io_exit(void)
{
    unregister_kprobe(&kp);
}

module_init(kprobe_io_init);
module_exit(kprobe_io_exit);
MODULE_LICENSE("GPL");
